#!/usr/bin/env python3
"""
无人机群仿真系统
================
使用Alpha混合技术将多架无人机叠加到背景视频上，生成逼真的无人机群飞行仿真视频
同时支持生成YOLO格式的机器学习训练标注数据
"""

import cv2
import numpy as np
import argparse
import random
from pathlib import Path
from tqdm import tqdm
import os


class DroneTrajectory:
    """无人机轨迹生成器 - 生成快速、平滑的线性运动轨迹"""
    
    def __init__(self, start_pos, end_pos, speed_pixels_per_frame):
        """初始化无人机轨迹
        
        Args:
            start_pos: 起始位置 [x, y]
            end_pos: 结束位置 [x, y] 
            speed_pixels_per_frame: 飞行速度（像素/帧）
        """
        self.start_pos = np.array(start_pos, dtype=float)  # 起始位置
        self.end_pos = np.array(end_pos, dtype=float)      # 结束位置
        
        # 根据距离和速度计算飞行持续时间（帧数）
        distance = np.linalg.norm(self.end_pos - self.start_pos)  # 计算欧几里得距离
        self.duration_frames = max(1, int(distance / speed_pixels_per_frame))  # 飞行总帧数
        
        # 计算方向向量，用于线性插值运动
        self.direction = (self.end_pos - self.start_pos) / self.duration_frames
    
    def get_position(self, frame):
        """获取指定帧数时无人机的位置
        
        Args:
            frame: 当前帧数
            
        Returns:
            当前位置坐标 [x, y]
        """
        if frame >= self.duration_frames:
            return self.end_pos  # 已到达终点
        
        # 使用线性插值计算当前位置：起始位置 + 方向向量 * 当前帧数
        return self.start_pos + (self.direction * frame)
    
    def is_finished(self, frame):
        """检查无人机是否已到达目的地
        
        Args:
            frame: 当前帧数
            
        Returns:
            bool: 是否已完成飞行
        """
        return frame >= self.duration_frames


class DroneSwarmSimulator:
    """无人机群仿真器主类 - 处理视频处理和无人机叠加的核心逻辑"""
    
    def __init__(self, background_video, drone_image, num_drones=10, gen_groundtruth=False):
        """初始化无人机群仿真器
        
        Args:
            background_video: 背景视频文件路径
            drone_image: 无人机图像文件路径
            num_drones: 同时存在的无人机数量
            gen_groundtruth: 是否生成机器学习训练用的标注数据
        """
        self.background_video = background_video    # 背景视频路径
        self.drone_image_path = drone_image        # 无人机图像路径
        self.num_drones = num_drones              # 无人机数量
        self.gen_groundtruth = gen_groundtruth    # 是否生成标注数据
        
        # 加载和准备多个无人机图像模型
        self.drone_images = []  # 存储所有无人机图像
        drone_paths = []        # 无人机图像文件路径列表
        
        # 在assets目录中查找所有无人机图像文件（drone*.png）
        assets_dir = Path(drone_image).parent
        for drone_file in assets_dir.glob("drone*.png"):
            drone_paths.append(str(drone_file))
        
        # 排序确保一致的加载顺序
        drone_paths.sort()
        
        # 如果没有找到drone*.png文件，使用指定的drone_image作为后备
        if not drone_paths:
            drone_paths = [drone_image]
        
        # 加载所有无人机图像
        for drone_path in drone_paths:
            # 使用IMREAD_UNCHANGED保持Alpha通道
            drone_img = cv2.imread(drone_path, cv2.IMREAD_UNCHANGED)
            if drone_img is None:
                print(f"警告: 无法加载无人机图像: {drone_path}")
                continue
            
            # 确保无人机图像具有Alpha通道（透明度）
            if drone_img.shape[2] == 3:  # 如果是RGB图像（缺少Alpha通道）
                # 添加Alpha通道，设置为完全不透明（255）
                alpha = np.ones((drone_img.shape[0], drone_img.shape[1], 1), dtype=drone_img.dtype) * 255
                drone_img = np.concatenate([drone_img, alpha], axis=2)  # 合并为RGBA
            
            self.drone_images.append(drone_img)
        
        # 检查是否成功加载了无人机图像
        if not self.drone_images:
            raise ValueError(f"无法从以下路径加载任何无人机图像: {drone_paths}")
        
        print(f"已加载 {len(self.drone_images)} 个无人机模型: {[Path(p).name for p in drone_paths]}")
        
        # 保持对第一个无人机图像的引用（向后兼容）
        self.drone_img = self.drone_images[0]
        
        # 加载背景视频
        self.cap = cv2.VideoCapture(background_video)
        if not self.cap.isOpened():
            raise ValueError(f"无法打开视频文件: {background_video}")
        
        # 获取视频属性
        self.fps = int(self.cap.get(cv2.CAP_PROP_FPS))              # 帧率
        self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))    # 视频宽度
        self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 视频高度
        self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))  # 总帧数
        
        print(f"视频信息: {self.width}x{self.height}, {self.fps} FPS, {self.total_frames} 帧")
        
        # 初始化无人机列表，用于连续的无人机流
        self.drones = []
        
        # 创建初始的无人机群
        for _ in range(self.num_drones):
            self.drones.append(self._create_new_drone(0))
    
    def _generate_random_edge_position(self, side):
        """在指定边缘生成随机位置
        
        Args:
            side: 边缘位置 ('left', 'right', 'top', 'bottom')
            
        Returns:
            [x, y]: 边缘上的随机坐标
        """
        if side == 'left':      # 左边缘
            return [0, random.randint(0, self.height)]
        elif side == 'right':   # 右边缘
            return [self.width, random.randint(0, self.height)]
        elif side == 'top':     # 上边缘
            return [random.randint(0, self.width), 0]
        else:  # bottom        # 下边缘
            return [random.randint(0, self.width), self.height]
    
    def _create_new_drone(self, current_frame):
        """创建一个具有随机轨迹和属性的新无人机
        
        Args:
            current_frame: 当前帧数
            
        Returns:
            dict: 包含无人机所有属性的字典
        """
        # 选择随机的起始和结束边缘（确保它们不同）
        edges = ['left', 'right', 'top', 'bottom']  # 四个边缘
        start_side = random.choice(edges)           # 随机选择起始边缘
        end_side = random.choice([e for e in edges if e != start_side])  # 选择不同的结束边缘
        
        # 生成起始和结束位置
        start_pos = self._generate_random_edge_position(start_side)
        end_pos = self._generate_random_edge_position(end_side)
        
        # 设置飞行速度（像素/帧），用于快速穿越屏幕
        speed = random.uniform(3, 8)  # 3-8像素每帧的随机速度
        
        # 创建轨迹对象
        trajectory = DroneTrajectory(start_pos, end_pos, speed)
        
        # 无人机颜色选项（用于着色变换）
        colors = [
            (255, 255, 255),  # 白色
            (0, 0, 0),        # 黑色
            (0, 0, 255),      # 红色
            (0, 255, 0),      # 绿色
            (255, 0, 0),      # 蓝色
            (0, 255, 255),    # 黄色
            (255, 0, 255),    # 品红色
            (255, 128, 0),    # 橙色
            (128, 0, 255),    # 紫色
            (0, 128, 255)     # 浅蓝色
        ]
        
        # 返回包含无人机所有属性的字典
        return {
            'trajectory': trajectory,                                           # 飞行轨迹
            'start_frame': current_frame,                                      # 开始帧
            'scale': random.uniform(0.015, 0.075),                           # 缩放比例（1.5%-7.5%）
            'rotation': random.uniform(0, 30),                                # 旋转角度（0-30度）
            'flip_horizontal': random.choice([True, False]),                  # 是否水平翻转
            'color': random.choice(colors),                                   # 颜色着色
            'alpha_multiplier': random.uniform(0.85, 1.0),                   # 透明度乘数
            'drone_image_index': random.randint(0, len(self.drone_images) - 1)  # 无人机模型索引
        }
    
    def _transform_drone(self, drone_img, scale, rotation, flip_horizontal, color):
        """对无人机图像进行变换：缩放、旋转、翻转和着色
        
        Args:
            drone_img: 原始无人机图像（RGBA格式）
            scale: 缩放比例
            rotation: 旋转角度
            flip_horizontal: 是否水平翻转
            color: 着色颜色 (B, G, R)
            
        Returns:
            变换后的无人机图像
        """
        # 1. 缩放图像
        h, w = drone_img.shape[:2]                    # 获取原始尺寸
        new_h, new_w = int(h * scale), int(w * scale) # 计算新尺寸
        resized = cv2.resize(drone_img, (new_w, new_h), interpolation=cv2.INTER_AREA)
        
        # 2. 应用颜色着色
        if resized.shape[2] == 4:  # 如果是RGBA图像
            # 分离RGB和Alpha通道
            rgb = resized[:, :, :3].astype(float)  # RGB通道
            alpha = resized[:, :, 3:4]             # Alpha通道
            
            # 仅对RGB通道应用颜色着色
            # 将现有颜色标准化并应用着色
            tinted_rgb = rgb * np.array(color) / 255.0
            tinted_rgb = np.clip(tinted_rgb, 0, 255).astype(np.uint8)  # 限制在0-255范围
            
            # 重新组合RGB和Alpha通道
            resized = np.concatenate([tinted_rgb, alpha], axis=2)
        
        # 3. 水平翻转（如果需要）
        if flip_horizontal:
            resized = cv2.flip(resized, 1)  # 1 = 水平翻转
        
        # 4. 旋转图像
        center = (new_w // 2, new_h // 2)  # 旋转中心点
        rotation_matrix = cv2.getRotationMatrix2D(center, rotation, 1.0)  # 创建旋转矩阵
        transformed = cv2.warpAffine(resized, rotation_matrix, (new_w, new_h), 
                                    flags=cv2.INTER_LINEAR,                    # 线性插值
                                    borderMode=cv2.BORDER_CONSTANT,           # 常数边界
                                    borderValue=(0, 0, 0, 0))                 # 透明边界
        
        return transformed
    
    def _blend_drone_onto_frame(self, frame, drone_img, position, alpha_multiplier=1.0):
        """使用Alpha混合将无人机叠加到视频帧上
        
        Args:
            frame: 背景视频帧
            drone_img: 无人机图像（RGBA格式）
            position: 无人机中心位置 [x, y]
            alpha_multiplier: 透明度乘数
            
        Returns:
            叠加后的视频帧
        """
        drone_h, drone_w = drone_img.shape[:2]  # 获取无人机图像尺寸
        # 计算无人机图像左上角位置（以中心点为基准）
        x, y = int(position[0] - drone_w // 2), int(position[1] - drone_h // 2)
        
        # 边界检查：如果无人机完全在画面外，直接返回原帧
        if x + drone_w <= 0 or y + drone_h <= 0 or x >= self.width or y >= self.height:
            return frame
        
        # 计算重叠区域（在视频帧中的有效区域）
        x1 = max(0, x)                          # 左边界
        y1 = max(0, y)                          # 上边界
        x2 = min(self.width, x + drone_w)       # 右边界
        y2 = min(self.height, y + drone_h)      # 下边界
        
        # 计算无人机图像中对应的区域
        dx1 = max(0, -x)                        # 无人机图像左偏移
        dy1 = max(0, -y)                        # 无人机图像上偏移
        dx2 = dx1 + (x2 - x1)                   # 无人机图像右边界
        dy2 = dy1 + (y2 - y1)                   # 无人机图像下边界
        
        # 如果没有有效的重叠区域，返回原帧
        if x2 <= x1 or y2 <= y1:
            return frame
        
        # 提取重叠区域
        frame_region = frame[y1:y2, x1:x2]           # 背景帧的重叠区域
        drone_region = drone_img[dy1:dy2, dx1:dx2]   # 无人机图像的重叠区域
        
        # Alpha混合处理
        if drone_region.shape[2] == 4:  # 如果无人机图像有Alpha通道
            # 提取Alpha通道并应用透明度乘数
            alpha = drone_region[:, :, 3:4].astype(float) / 255.0 * alpha_multiplier
            drone_rgb = drone_region[:, :, :3]  # 无人机的RGB通道
            
            # Alpha混合公式：result = background * (1-alpha) + foreground * alpha
            blended = frame_region * (1 - alpha) + drone_rgb * alpha
            frame[y1:y2, x1:x2] = blended.astype(np.uint8)  # 更新背景帧
        
        return frame
    
    def _get_drone_bounding_box(self, position, drone_img):
        """计算无人机在指定位置的边界框
        
        Args:
            position: 无人机中心位置 [x, y]
            drone_img: 无人机图像
            
        Returns:
            边界框坐标 (x1, y1, x2, y2) 或 None
        """
        drone_h, drone_w = drone_img.shape[:2]  # 获取无人机图像尺寸
        x_center = position[0]  # 中心X坐标
        y_center = position[1]  # 中心Y坐标
        
        # 计算边界框的四个角点
        x1 = max(0, x_center - drone_w // 2)           # 左边界
        y1 = max(0, y_center - drone_h // 2)           # 上边界
        x2 = min(self.width, x_center + drone_w // 2)  # 右边界
        y2 = min(self.height, y_center + drone_h // 2) # 下边界
        
        # 只有当边界框在画面内可见时才返回
        if x2 > x1 and y2 > y1:
            return (x1, y1, x2, y2)
        return None
    
    def _bbox_to_yolo_format(self, bbox):
        """将边界框转换为YOLO格式的标注
        
        Args:
            bbox: 边界框坐标 (x1, y1, x2, y2)
            
        Returns:
            YOLO格式字符串: "class_id x_center y_center width height" (归一化坐标)
        """
        if bbox is None:
            return None
        
        x1, y1, x2, y2 = bbox
        # 计算归一化的中心坐标和尺寸
        x_center = (x1 + x2) / 2.0 / self.width   # 归一化中心X坐标
        y_center = (y1 + y2) / 2.0 / self.height  # 归一化中心Y坐标
        width = (x2 - x1) / self.width            # 归一化宽度
        height = (y2 - y1) / self.height          # 归一化高度
        
        # 类别ID 0 表示UAV（无人机）
        return f"0 {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}"
    
    def _draw_bounding_box(self, frame, bbox, label="UAV"):
        """在视频帧上绘制绿色边界框和标签
        
        Args:
            frame: 视频帧
            bbox: 边界框坐标 (x1, y1, x2, y2)
            label: 标签文本
            
        Returns:
            绘制了边界框的视频帧
        """
        if bbox is None:
            return frame
        
        x1, y1, x2, y2 = bbox
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)  # 转换为整数坐标
        
        # 绘制绿色边界框
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        
        # 添加标签
        label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)[0]
        # 绘制标签背景
        cv2.rectangle(frame, (x1, y1 - label_size[1] - 5), (x1 + label_size[0], y1), (0, 255, 0), -1)
        # 绘制标签文本
        cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
        
        return frame
    
    def simulate(self, output_path="output_swarm.mp4"):
        """运行仿真并保存输出视频
        
        Args:
            output_path: 输出视频文件路径
        """
        # 创建目录结构: output_path/<vid_name>/<images/ | vid_name.mp4 | bbox_vid_name.mp4 | labels/>
        output_path_obj = Path(output_path)
        vid_name = output_path_obj.stem  # 获取不带扩展名的文件名
        
        # 创建主视频目录
        vid_dir = output_path_obj.parent / vid_name
        vid_dir.mkdir(parents=True, exist_ok=True)
        
        # 创建子目录
        images_dir = vid_dir / "images"  # 单帧图像目录
        images_dir.mkdir(exist_ok=True)
        
        # 更新输出路径
        video_output_path = vid_dir / f"{vid_name}.mp4"
        
        # 设置视频写入器
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # MP4编码器
        out = cv2.VideoWriter(str(video_output_path), fourcc, self.fps, (self.width, self.height))
        
        # 如果启用了标注数据生成，设置相关组件
        bbox_out = None
        labels_dir = None
        if self.gen_groundtruth:
            # 在同一目录中创建带标注的视频
            bbox_video_path = vid_dir / f"bbox_{vid_name}.mp4"
            bbox_out = cv2.VideoWriter(str(bbox_video_path), fourcc, self.fps, (self.width, self.height))
            
            # 创建标注文件目录
            labels_dir = vid_dir / "labels"
            labels_dir.mkdir(exist_ok=True)
            
            print(f"已启用标注数据生成。带标注视频: {bbox_video_path}")
            print(f"YOLO标注文件将保存在: {labels_dir}")
        
        # 使用tqdm进度条处理视频帧
        with tqdm(total=self.total_frames, desc="处理视频帧", unit="帧") as pbar:
            frame_num = 0
            
            while True:
                ret, frame = self.cap.read()  # 读取视频帧
                if not ret:
                    break  # 视频结束
                
                # 更新无人机位置并移除已完成飞行的无人机
                active_drones = []    # 活跃的无人机列表
                frame_bboxes = []     # 当前帧的边界框列表
                
                for drone in self.drones:
                    relative_frame = frame_num - drone['start_frame']  # 计算相对帧数
                    
                    if not drone['trajectory'].is_finished(relative_frame):
                        # 无人机仍在飞行中
                        position = drone['trajectory'].get_position(relative_frame)  # 获取当前位置
                        
                        # 变换无人机图像（使用该无人机的特定模型）
                        selected_drone_img = self.drone_images[drone['drone_image_index']]
                        transformed_drone = self._transform_drone(
                            selected_drone_img, drone['scale'], drone['rotation'], 
                            drone['flip_horizontal'], drone['color']
                        )
                        
                        # 如果启用了标注数据生成，计算边界框
                        if self.gen_groundtruth:
                            bbox = self._get_drone_bounding_box(position, transformed_drone)
                            if bbox:
                                frame_bboxes.append(bbox)
                        
                        # 将无人机叠加到视频帧上
                        frame = self._blend_drone_onto_frame(
                            frame, transformed_drone, position, drone['alpha_multiplier']
                        )
                        
                        active_drones.append(drone)
                
                # 用新无人机替换已完成飞行的无人机
                while len(active_drones) < self.num_drones:
                    # 添加一些随机性到生成时机
                    if random.random() < 0.3:  # 每帧30%的概率生成新无人机
                        active_drones.append(self._create_new_drone(frame_num))
                
                self.drones = active_drones  # 更新无人机列表
                
                # 保存单帧图像到images目录
                frame_filename = images_dir / f"{frame_num:06d}.jpg"
                cv2.imwrite(str(frame_filename), frame)
                
                # 如果启用了标注数据生成，写入标注数据
                if self.gen_groundtruth:
                    # 写入YOLO格式的标注文件
                    label_file = labels_dir / f"{frame_num:06d}.txt"
                    with open(label_file, 'w') as f:
                        for bbox in frame_bboxes:
                            yolo_annotation = self._bbox_to_yolo_format(bbox)
                            if yolo_annotation:
                                f.write(yolo_annotation + '\n')
                    
                    # 创建带边界框的标注视频帧
                    if bbox_out:
                        annotated_frame = frame.copy()
                        for bbox in frame_bboxes:
                            annotated_frame = self._draw_bounding_box(annotated_frame, bbox, "UAV")
                        bbox_out.write(annotated_frame)
                
                out.write(frame)  # 写入主视频
                frame_num += 1
                pbar.update(1)    # 更新进度条
        
        # 清理资源
        self.cap.release()    # 释放视频捕获对象
        out.release()         # 释放视频写入器
        if bbox_out:
            bbox_out.release()  # 释放标注视频写入器
        
        print(f"仿真完成！输出已保存在: {vid_dir}")
        print(f"  - 主视频: {video_output_path}")
        print(f"  - 单帧图像: {images_dir}")
        if self.gen_groundtruth:
            print(f"  - 标注数据: {labels_dir}")
            print(f"  - 带标注视频: {bbox_video_path}")


def main():
    """主函数 - 解析命令行参数并运行仿真"""
    parser = argparse.ArgumentParser(description="无人机群视频仿真器")
    parser.add_argument("--background", default="assets/background.mp4", 
                       help="背景视频文件 (默认: assets/background.mp4)")
    parser.add_argument("--drone", default="assets/drone.png",
                       help="无人机PNG图像文件 (默认: assets/drone.png)")
    parser.add_argument("--num-drones", type=int, default=10,
                       help="无人机群中的无人机数量 (默认: 10)")
    parser.add_argument("--output", default="outputs/output_swarm.mp4",
                       help="输出视频文件 (默认: outputs/output_swarm.mp4)")
    parser.add_argument("--gen_groundtruth", action="store_true",
                       help="生成YOLO格式的标注数据和带标注的验证视频")
    
    args = parser.parse_args()
    
    # 检查必需文件是否存在
    if not Path(args.background).exists():
        print(f"错误: 背景视频文件 '{args.background}' 未找到!")
        print("请将背景视频放在 'assets/background.mp4' 或使用 --background 指定路径")
        return 1
    
    if not Path(args.drone).exists():
        print(f"错误: 无人机图像文件 '{args.drone}' 未找到!")
        print("请将无人机PNG图像放在 'assets/drone.png' 或使用 --drone 指定路径")
        return 1
    
    try:
        # 确保输出目录存在
        output_dir = Path(args.output).parent
        output_dir.mkdir(exist_ok=True)
        
        # 创建仿真器并运行
        simulator = DroneSwarmSimulator(args.background, args.drone, args.num_drones, args.gen_groundtruth)
        simulator.simulate(args.output)
        return 0
    
    except Exception as e:
        print(f"错误: {e}")
        return 1


if __name__ == "__main__":
    exit(main())